""" Get the price computation values that speed up the code

"""
import pandas as pd
import numpy as np
import json
try:
    from src.interpolation.splines import UCGrid
except:
    from ..interpolation.splines import UCGrid
import sys
sys.path.append('./')

from src.models_new import prices, first_stage_utils

#%% CONFIGURE ---------------------------------------------------------------------------
polynomial_order = 2
time_periods = ['month', 'fortnight']

for i in ['', '_long']:
    for t in time_periods:
        #%% IMPORT DATA ---------------------------------------------------------------------
        if t == 'month':
            beta = 0.99
        elif t == 'fortnight':
            beta = 0.995

        df_contracts = pd.read_csv(f'./data_py/processed/contracts_final{i}.csv', index_col=[0])
        if t == 'fortnight':
            df_contracts['g'] = df_contracts['g_fortnight']
            df_contracts['n_l'] = df_contracts['n_l_fortnight']
            df_contracts['n_m'] = df_contracts['n_m_fortnight']
            df_contracts['n_h'] = df_contracts['n_h_fortnight']
        r = pd.read_csv(f'./models/first_stage/r_{t}.csv', index_col=[0])
        const = pd.read_csv(f'./models/first_stage/const_{t}.csv', index_col=[0])
        sigma = pd.read_csv(f'./models/first_stage/sigma_{t}.csv', index_col=[0])
        seeds = list(range(200))
        df_extensions_with_mri = pd.read_csv(
            f'./models/first_stage/first_stage_extensions_with_mri_{t}.csv', index_col=[0])

        search_grid_params_by_spec = dict()
        search_grid_by_spec = dict()
        search_value_by_spec = dict()
        prob_extend_by_spec = dict()
        for spec in ['low', 'mid', 'high']:
            search_grid_params_by_spec[spec] = np.load(
                f'./models/value_search/search_grid_{spec}_{t}.npy')
            search_grid_by_spec[spec] = UCGrid(
                tuple(search_grid_params_by_spec[spec][0, :]),
                tuple(search_grid_params_by_spec[spec][1, :]),
                tuple(search_grid_params_by_spec[spec][2, :]),
                tuple(search_grid_params_by_spec[spec][3, :])
            )
            search_value_by_spec[spec] = np.load(
                f'./models/value_search/search_value_{spec}_{t}.npy')
            prob_extend_by_spec[spec] = first_stage_utils.build_prob_extension_predictor_with_mri(
                coefs=df_extensions_with_mri.loc[spec].values)


        #%% GET THE PRICE-MATCH MOMENTS -----------------------------------------------------
        params_prices = {
            'm_0': 1.0,
            'm_1': 1.0,
            'm_2': 1.0,
            'rho_0': 1.0,
            'rho_1': 1.0,
            'rho_2': 1.0,
            'rho_3': 1.0,
            'delta': 0.0,  # NEED to input delta here (the other params do not make a difference)
            'beta': beta
        }

        price_match_values_by_spec = prices.build_fast_prices(
            df_contracts, search_grid_by_spec, search_value_by_spec, const,
            r, sigma, prob_extend_by_spec, params_prices, seeds)

        for spec in ['low', 'mid', 'high']:
            for k in [0, 1, 2, 3]:
                with open(f"./models/price_match/price_match_values_{spec}_{k}_{t}{i}.json", "w") as outfile:
                    json.dump(list(price_match_values_by_spec[spec][k]), outfile)
